import numpy as np
from typing import Dict, List, Tuple
import os
import logging
from datetime import datetime
from physics import PhysicsModel, PhaseType
from parallel_compute import (
    calculate_mass_flux_array,
    calculate_momentum_flux_array,
    calculate_wall_friction_array,
    calculate_interfacial_friction_array,
    calculate_lateral_flow_resistance_array,
    calculate_turbulent_mixing_array,
    solve_momentum_equation_parallel,
    solve_continuity_equation_parallel
)

class Solver:
    """两相流求解器类（使用并行计算）"""
    def __init__(self, grid, physics_model):
        self.grid = grid
        self.physics = physics_model
        
        # 求解器参数
        self.max_iterations = 1000
        self.convergence_tolerance = 1e-6
        self.relaxation_factors = {
            'pressure': 0.3,
            'velocity': 0.7,
            'void_fraction': 0.3
        }
        
        # CFL条件相关
        self.max_cfl = 0.8
        self.min_dt = 1e-6
        self.max_dt = 0.01
        
        # 时间步长
        self.dt = self.min_dt
        
        # 设置日志记录
        self._setup_logging()
        
        # 设置初始条件
        self.initialize_fields()
        
    def _setup_logging(self):
        """设置日志记录"""
        # 创建results文件夹（如果不存在）
        results_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'results')
        if not os.path.exists(results_dir):
            os.makedirs(results_dir)
        
        # 创建日志文件名（使用时间戳）
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = os.path.join(results_dir, f'simulation_{timestamp}.log')
        
        # 配置日志记录器
        self.logger = logging.getLogger('TwoFluidSCA')
        self.logger.setLevel(logging.INFO)
        
        # 文件处理器
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(logging.INFO)
        
        # 控制台处理器
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        
        # 创建格式化器
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)
        
        # 添加处理器
        self.logger.addHandler(file_handler)
        self.logger.addHandler(console_handler)
        
        self.logger.info(f"日志文件创建于: {log_file}")
        
    def log_and_print(self, message: str, level: str = 'info'):
        """同时记录日志并打印到控制台
        
        Args:
            message: 要记录的消息
            level: 日志级别 ('info', 'warning', 'error', 'debug')
        """
        if level == 'info':
            self.logger.info(message)
        elif level == 'warning':
            self.logger.warning(message)
        elif level == 'error':
            self.logger.error(message)
        elif level == 'debug':
            self.logger.debug(message)
        
    def initialize_fields(self):
        """初始化所有场量"""
        self.log_and_print("\n初始化计算场...")
        for subchannel in self.grid.subchannels:
            n_axial = len(subchannel.alpha)
            
            # 1. 压力场初始化：线性分布（常压，略高于大气压）
            P_in = 1.2e5  # 入口压力 (1.2 bar)
            P_out = 1.0e5  # 出口压力 (1.0 bar)
            subchannel.pressure = np.linspace(P_in, P_out, n_axial)
            subchannel.previous_pressure = subchannel.pressure.copy()
            
            # 2. 空泡份额初始化：入口值延拓（减小初始空泡份额以提高稳定性）
            alpha_in = 0.1  # 入口空泡份额（降低初始值）
            subchannel.alpha = np.full(n_axial, alpha_in)
            subchannel.previous_alpha = subchannel.alpha.copy()
            
            # 3. 速度场初始化（降低初始速度以提高稳定性）
            v_l_in = 0.15  # 入口液相速度
            v_g_in = 0.20  # 入口气相速度（略快于液相）
            
            # 初始化速度场（包括轴向和横向分量）
            subchannel.velocity = {
                'liquid': {
                    'axial': np.full(n_axial, v_l_in),
                    'lateral': np.zeros(n_axial)
                },
                'vapor': {
                    'axial': np.full(n_axial, v_g_in),
                    'lateral': np.zeros(n_axial)
                }
            }
            
            # 初始化previous_velocity
            subchannel.previous_velocity = {
                'liquid': {
                    'axial': subchannel.velocity['liquid']['axial'].copy(),
                    'lateral': subchannel.velocity['liquid']['lateral'].copy()
                },
                'vapor': {
                    'axial': subchannel.velocity['vapor']['axial'].copy(),
                    'lateral': subchannel.velocity['vapor']['lateral'].copy()
                }
            }
            
            # 初始化质量通量
            subchannel.previous_mass_flux = {
                'vapor': np.zeros(n_axial),
                'liquid': np.zeros(n_axial)
            }
            
            # 初始化漂移通量
            subchannel.previous_drift_flux = np.zeros(n_axial)
            
        self.log_and_print("场初始化完成：")
        self.log_and_print(f"- 压力范围：{P_in/1e5:.2f} bar - {P_out/1e5:.2f} bar")
        self.log_and_print(f"- 入口空泡份额：{alpha_in:.3f}")
        self.log_and_print(f"- 入口速度：液相 {v_l_in:.2f} m/s, 气相 {v_g_in:.2f} m/s")
        
    def calculate_time_step(self) -> float:
        """计算自适应时间步长"""
        # 计算最大速度
        max_velocity = 1e-6  # 设置一个最小速度以避免除零
        for subchannel in self.grid.subchannels:
            for phase in ['liquid', 'vapor']:
                max_velocity = max(max_velocity,
                                 np.max(np.abs(subchannel.velocity[phase]['axial'])),
                                 np.max(np.abs(subchannel.velocity[phase]['lateral'])))
        
        # 基于CFL条件计算时间步长
        dt_cfl = self.max_cfl * min(self.grid.dx, self.grid.dz) / max_velocity
        
        # 限制时间步长在允许范围内
        self.dt = np.clip(dt_cfl, self.min_dt, self.max_dt)
        
        return self.dt
    
    def predict_velocity(self, subchannel, dt: float):
        """速度预测步骤"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 使用动量方程预测速度
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                # 计算压力梯度
                dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                
                # 预测速度
                subchannel.velocity[phase]['axial'][i] = (
                    subchannel.previous_velocity[phase]['axial'][i] -
                    dt * dp_dz / props['density']
                )
    
    def solve_pressure_correction(self, subchannel):
        """求解压力修正方程"""
        n_axial = len(subchannel.alpha)
        
        # 构建压力修正方程系数矩阵
        A = np.zeros((n_axial, n_axial))
        b = np.zeros(n_axial)
        
        # 设置入口边界条件（固定压力）
        A[0, 0] = 1.0
        b[0] = 0.0  # 入口压力不需要修正
        
        # 设置出口边界条件（固定压力）
        A[-1, -1] = 1.0
        b[-1] = 0.0  # 出口压力不需要修正
        
        # 构建内部节点的方程
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算密度项（添加保护）
            rho_g = max(props_g['density'], 1e-6)
            rho_l = max(props_l['density'], 1e-6)
            
            # 计算速度项
            v_g = subchannel.velocity['vapor']['axial'][i]
            v_l = subchannel.velocity['liquid']['axial'][i]
            
            # 计算系数（考虑密度和空泡份额的影响，添加保护）
            alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
            coeff = (alpha / rho_g + (1 - alpha) / rho_l)
            
            # 添加数值稳定项
            stability_factor = 1e-8
            
            # 设置系数矩阵
            A[i, i-1] = -coeff / (self.grid.dz ** 2)
            A[i, i] = 2 * coeff / (self.grid.dz ** 2) + stability_factor
            A[i, i+1] = -coeff / (self.grid.dz ** 2)
            
            # 计算质量不平衡作为源项
            G_g = alpha * rho_g * v_g
            G_l = (1 - alpha) * rho_l * v_l
            
            # 计算质量流量的空间导数（使用中心差分）
            alpha_prev = np.clip(subchannel.alpha[i-1], 1e-6, 1-1e-6)
            v_g_prev = subchannel.velocity['vapor']['axial'][i-1]
            v_l_prev = subchannel.velocity['liquid']['axial'][i-1]
            
            G_g_prev = alpha_prev * rho_g * v_g_prev
            G_l_prev = (1 - alpha_prev) * rho_l * v_l_prev
            
            dG_g_dz = (G_g - G_g_prev) / self.grid.dz
            dG_l_dz = (G_l - G_l_prev) / self.grid.dz
            
            b[i] = -(dG_g_dz + dG_l_dz)
        
        try:
            # 添加对角优势以提高稳定性
            for i in range(n_axial):
                A[i,i] = max(A[i,i], 1.1 * np.sum(np.abs(A[i,:])))
            
            # 使用直接求解器求解压力修正方程
            pressure_correction = np.linalg.solve(A, b)
            
            # 检查解的有效性
            if not np.all(np.isfinite(pressure_correction)):
                raise ValueError("压力修正解包含无效值")
            
            # 限制压力修正的幅度
            max_correction = 0.1 * np.mean(subchannel.pressure)
            pressure_correction = np.clip(pressure_correction, -max_correction, max_correction)
            
            # 应用压力修正（带松弛因子）
            subchannel.pressure += self.relaxation_factors['pressure'] * pressure_correction
            
            # 确保压力始终在物理范围内
            subchannel.pressure = np.maximum(subchannel.pressure, 1e5)  # 最小压力为1个大气压
            
        except (np.linalg.LinAlgError, ValueError) as e:
            print(f"警告：压力修正方程求解失败: {str(e)}")
            # 使用简单的压力梯度估计
            dp = (subchannel.pressure[-1] - subchannel.pressure[0]) / (n_axial - 1)
            subchannel.pressure = np.linspace(subchannel.pressure[0], subchannel.pressure[-1], n_axial)
    
    def correct_velocity(self, subchannel):
        """基于压力修正更新速度场"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            for phase in ['vapor', 'liquid']:
                # 计算速度修正
                dp_correction = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                velocity_correction = -dp_correction * self.relaxation_factors['velocity']
                
                # 更新速度
                subchannel.velocity[phase]['axial'][i] += velocity_correction
    
    def solve_subchannel_parallel(self, subchannel, dt: float):
        """并行求解单个子通道，包含横向流动和空隙漂移"""
        n_axial = len(subchannel.alpha)
        
        # 获取相邻子通道信息
        adjacent_channels = self.grid.get_adjacent_channels(subchannel.id)
        
        # 获取几何参数
        geom = subchannel.get_geometry_parameters()
        
        # 准备数组数据
        alpha = subchannel.alpha
        pressure = subchannel.pressure
        v_g_axial = subchannel.velocity['vapor']['axial']
        v_l_axial = subchannel.velocity['liquid']['axial']
        v_g_lateral = subchannel.velocity['vapor']['lateral']
        v_l_lateral = subchannel.velocity['liquid']['lateral']
        
        # 获取物性
        props_g = self.physics.calculate_properties(pressure[0], PhaseType.VAPOR)
        props_l = self.physics.calculate_properties(pressure[0], PhaseType.LIQUID)
        rho_g = props_g['density']
        rho_l = props_l['density']
        
        # 计算空隙漂移速度
        v_gj, C0 = self.physics.calculate_drift_velocity(rho_g, rho_l)
        
        # 修正气相速度（考虑漂移速度）
        j = alpha * v_g_axial + (1 - alpha) * v_l_axial  # 混合体积通量
        v_g_axial = C0 * j + v_gj
        
        # 轴向动量方程
        for i in range(1, n_axial-1):
            # 计算轴向压力梯度
            dp_dz = (pressure[i+1] - pressure[i-1]) / (2 * self.grid.dz)
            
            # 计算轴向壁面摩擦力
            F_w_g = self.physics.calculate_wall_friction(
                PhaseType.VAPOR, rho_g, v_g_axial[i],
                geom['hydraulic_diameter'], props_g['viscosity']
            )
            F_w_l = self.physics.calculate_wall_friction(
                PhaseType.LIQUID, rho_l, v_l_axial[i],
                geom['hydraulic_diameter'], props_l['viscosity']
            )
            
            # 计算界面摩擦力
            F_i = self.physics.calculate_interfacial_friction(
                alpha[i], rho_g, rho_l, v_g_axial[i], v_l_axial[i],
                geom['hydraulic_diameter']
            )
            
            # 横向动量交换（遍历相邻通道）
            M_lat_g_total = 0
            M_lat_l_total = 0
            for adj_ch in adjacent_channels:
                # 获取相邻通道数据
                v_g_adj = adj_ch.velocity['vapor']['axial'][i]
                v_l_adj = adj_ch.velocity['liquid']['axial'][i]
                
                # 获取连接参数
                conn_params = subchannel.get_connection_parameters(adj_ch)
                
                # 计算横向动量交换
                M_lat_g, M_lat_l = self.physics.calculate_lateral_momentum_exchange(
                    alpha[i], rho_g, rho_l,
                    v_g_axial[i], v_l_axial[i],
                    v_g_adj, v_l_adj,
                    conn_params['gap_length'],
                    conn_params['gap_width']
                )
                
                # 计算湍流交混
                w_t = self.physics.calculate_turbulent_mixing(
                    self.physics.BETA,
                    conn_params['gap_width'],
                    rho_g, rho_l,
                    v_g_axial[i], v_l_axial[i]
                )
                
                M_t = self.physics.calculate_turbulent_mixing_momentum(
                    w_t, v_g_axial[i], v_g_adj
                )
                
                M_lat_g_total += M_lat_g + M_t
                M_lat_l_total += M_lat_l + M_t
            
            # 更新轴向速度
            v_g_axial[i] += dt * (
                -dp_dz/rho_g + F_w_g + F_i/alpha[i] + 
                M_lat_g_total/(rho_g * geom['area']) - self.physics.g
            )
            
            v_l_axial[i] += dt * (
                -dp_dz/rho_l + F_w_l - F_i/(1-alpha[i]) + 
                M_lat_l_total/(rho_l * geom['area']) - self.physics.g
            )
        
        # 更新横向速度（使用横向压力梯度）
        for adj_ch in adjacent_channels:
            conn_params = subchannel.get_connection_parameters(adj_ch)
            dp_dx = (adj_ch.pressure - pressure) / conn_params['centroid_distance']
            v_g_lateral += dt * (-dp_dx/rho_g)
            v_l_lateral += dt * (-dp_dx/rho_l)
        
        # 更新空泡份额（考虑漂移速度的影响）
        j_drift = alpha * (1 - alpha) * v_gj
        alpha_new = alpha + dt * (
            -np.gradient(alpha * v_g_axial, self.grid.dz) +
            np.gradient(j_drift, self.grid.dz)
        )
        
        # 限制空泡份额在物理范围内
        alpha_new = np.clip(alpha_new, 0.0, 1.0)
        
        # 更新结果
        subchannel.alpha = alpha_new
        subchannel.velocity['vapor']['axial'] = v_g_axial
        subchannel.velocity['liquid']['axial'] = v_l_axial
        subchannel.velocity['vapor']['lateral'] = v_g_lateral
        subchannel.velocity['liquid']['lateral'] = v_l_lateral
        
        # 计算残差
        residuals = {
            'velocity_g': np.max(np.abs(v_g_axial - subchannel.previous_velocity['vapor']['axial'])),
            'velocity_l': np.max(np.abs(v_l_axial - subchannel.previous_velocity['liquid']['axial'])),
            'alpha': np.max(np.abs(alpha_new - subchannel.previous_alpha))
        }
        
        return residuals
    
    def solve_time_step(self):
        """求解一个时间步"""
        # 计算时间步长
        self.dt = self.calculate_time_step()
        self.log_and_print(f"\n时间步长: {self.dt:.3e} s")
        
        # 预计算物性
        P_min = min(ch.pressure.min() for ch in self.grid.subchannels)
        P_max = max(ch.pressure.max() for ch in self.grid.subchannels)
        self.log_and_print(f"压力范围: {P_min/1e6:.3f} MPa - {P_max/1e6:.3f} MPa")
        self.physics.precalculate_properties(P_min, P_max, n_points=10)
        
        # 保存当前状态
        for subchannel in self.grid.subchannels:
            subchannel.save_current_state()
        
        # 改进的PISO算法主循环
        self.log_and_print("\nPISO算法迭代:")
        for iter in range(self.max_iterations):
            max_residuals = {
                'velocity_g': 0.0,
                'velocity_l': 0.0,
                'alpha': 0.0,
                'pressure': 0.0
            }
            
            # 对每个子通道求解
            for subchannel in self.grid.subchannels:
                # 1. 使用改进的PISO算法
                self.solve_improved_piso(subchannel)
                
                # 2. 使用TVD格式更新空泡份额
                self.update_void_fraction_tvd(subchannel)
                
                # 3. 计算残差
                try:
                    # 计算特征速度和压力
                    v_g_char = max(np.max(np.abs(subchannel.velocity['vapor']['axial'])), 0.1)
                    v_l_char = max(np.max(np.abs(subchannel.velocity['liquid']['axial'])), 0.1)
                    p_char = max(np.mean(np.abs(subchannel.pressure)), 1e5)
                    
                    # 计算相对残差
                    residuals = {
                        'velocity_g': np.max(np.abs(subchannel.velocity['vapor']['axial'] - 
                                               subchannel.previous_velocity['vapor']['axial'])) / v_g_char,
                        'velocity_l': np.max(np.abs(subchannel.velocity['liquid']['axial'] - 
                                               subchannel.previous_velocity['liquid']['axial'])) / v_l_char,
                        'alpha': np.max(np.abs(subchannel.alpha - subchannel.previous_alpha)),
                        'pressure': np.max(np.abs(subchannel.pressure - subchannel.previous_pressure)) / p_char
                    }
                    
                    # 限制残差的范围
                    for key in residuals:
                        if np.isfinite(residuals[key]):
                            residuals[key] = min(residuals[key], 1.0)
                            max_residuals[key] = max(max_residuals[key], residuals[key])
                        else:
                            self.log_and_print(f"警告：{key}残差计算结果无效", level='warning')
                            
                except Exception as e:
                    self.log_and_print(f"警告：残差计算失败: {str(e)}", level='warning')
                    continue
            
            # 输出残差
            self.log_and_print(
                f"迭代 {iter:4d}: " + 
                f"气相速度残差={max_residuals['velocity_g']:.3e}, " +
                f"液相速度残差={max_residuals['velocity_l']:.3e}, " +
                f"空泡份额残差={max_residuals['alpha']:.3e}, " +
                f"压力残差={max_residuals['pressure']:.3e}"
            )
            
            # 检查收敛性
            if all(r < self.convergence_tolerance for r in max_residuals.values()):
                self.log_and_print(f"\n已收敛，总迭代次数: {iter + 1}")
                break
            elif iter == self.max_iterations - 1:
                self.log_and_print(f"\n警告：达到最大迭代次数 {self.max_iterations}，未收敛", level='warning')
        
        # 输出当前状态
        center_channel = self.grid.subchannels[0]  # 使用第一个通道作为示例
        self.log_and_print("\n当前状态:")
        
        try:
            self.log_and_print(f"入口压力: {center_channel.pressure[0]/1e6:.3f} MPa")
            self.log_and_print(f"出口压力: {center_channel.pressure[-1]/1e6:.3f} MPa")
        except:
            self.log_and_print("警告：压力值计算无效", level='warning')
        
        self.log_and_print(f"入口空泡份额: {center_channel.alpha[0]:.3f}")
        self.log_and_print(f"入口气相速度: {center_channel.velocity['vapor']['axial'][0]:.3f} m/s")
        self.log_and_print(f"入口液相速度: {center_channel.velocity['liquid']['axial'][0]:.3f} m/s")
        
        # 清空物性缓存，准备下一个时间步
        self.physics.clear_cache()
        
        # 更新时间
        self.grid.advance_time(self.dt)
    def solve_improved_piso(self, subchannel):
        """使用改进的PISO算法求解速度和压力场
    
        Args:
            subchannel: 子通道对象
        """
        # 第一步：速度预测
        self.predict_velocity_piso(subchannel)
        
        # 第二步：压力修正（使用Rhie-Chow插值）
        self.solve_pressure_correction_rhie_chow(subchannel)
        
        # 第三步：考虑界面动量交换的压力修正
        self.solve_pressure_correction_with_interfacial(subchannel)
        
        # 第四步：速度修正
        self.correct_velocity_with_interfacial(subchannel)

    def solve_pressure_correction_rhie_chow(self, subchannel):
        """使用Rhie-Chow插值求解压力修正方程"""
        # 获取物性参数
        n = len(subchannel.pressure)
        
        # 构建系数矩阵
        A = np.zeros((n, n))
        b = np.zeros(n)
        
        # 边界条件
        A[0, 0] = 1.0
        b[0] = 0.0  # 入口压力修正为0
        A[-1, -1] = 1.0
        b[-1] = 0.0  # 出口压力修正为0
        
        # 内部节点
        for i in range(1, n-1):
            try:
                # 获取物性（添加错误检查）
                props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
                props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
                
                # 确保密度为正值
                rho_g = max(props_g['density'], 1e-6)
                rho_l = max(props_l['density'], 1e-6)
                
                # 确保空泡份额在物理范围内
                alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
                
                # 计算系数
                ap_g = rho_g / self.dt
                ap_l = rho_l / self.dt
                
                # 添加数值稳定项
                stability_factor = 1e-8
                
                # 构建方程系数
                A[i, i-1] = -(alpha/rho_g + (1-alpha)/rho_l) / (self.grid.dz**2)
                A[i, i] = 2*(alpha/rho_g + (1-alpha)/rho_l) / (self.grid.dz**2) + stability_factor
                A[i, i+1] = -(alpha/rho_g + (1-alpha)/rho_l) / (self.grid.dz**2)
                
                # 计算源项（质量不平衡）
                mass_imbalance = self.calculate_mass_imbalance(subchannel, i)
                if np.isfinite(mass_imbalance):
                    b[i] = mass_imbalance
                else:
                    self.log_and_print(f"警告：节点 {i} 的质量不平衡计算结果无效", level='warning')
                    b[i] = 0.0
                    
            except Exception as e:
                self.log_and_print(f"警告：节点 {i} 的压力修正方程系数计算失败: {str(e)}", level='warning')
                # 使用简单的拉普拉斯算子作为后备
                A[i, i-1] = -1.0
                A[i, i] = 2.0
                A[i, i+1] = -1.0
                b[i] = 0.0
        
        try:
            # 增加对角优势以提高稳定性
            for i in range(n):
                A[i,i] = max(A[i,i], 1.1 * np.sum(np.abs(A[i,:])))
            
            # 求解压力修正方程
            dp = np.linalg.solve(A, b)
            
            # 检查解的有效性
            if not np.all(np.isfinite(dp)):
                raise ValueError("压力修正解包含无效值")
            
            # 限制修正幅度
            max_correction = 0.1 * np.mean(np.abs(subchannel.pressure))
            dp = np.clip(dp, -max_correction, max_correction)
            
            # 应用压力修正（使用较小的松弛因子）
            omega = 0.2  # 更小的松弛因子
            subchannel.pressure += omega * dp
            
            # 确保压力保持在物理范围内
            min_pressure = 1e5  # 最小压力（1 bar）
            subchannel.pressure = np.maximum(subchannel.pressure, min_pressure)
            
        except Exception as e:
            self.log_and_print(f"警告：压力修正方程求解失败: {str(e)}", level='warning')
            # 使用简单的压力梯度估计
            dp = (subchannel.pressure[-1] - subchannel.pressure[0]) / (n - 1)
            subchannel.pressure = np.linspace(subchannel.pressure[0], subchannel.pressure[-1], n)

    def calculate_mass_imbalance(self, subchannel, i):
        """计算质量不平衡
        
        Args:
            subchannel: 子通道对象
            i: 当前节点索引
            
        Returns:
            float: 质量不平衡
        """
        try:
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 确保密度为正值
            rho_g = max(props_g['density'], 1e-6)
            rho_l = max(props_l['density'], 1e-6)
            
            # 确保空泡份额在物理范围内
            alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
            alpha_prev = np.clip(subchannel.alpha[i-1], 1e-6, 1-1e-6)
            
            # 计算速度项（添加限制以避免过大的速度）
            v_g = np.clip(subchannel.velocity['vapor']['axial'][i], -10.0, 10.0)
            v_l = np.clip(subchannel.velocity['liquid']['axial'][i], -10.0, 10.0)
            v_g_prev = np.clip(subchannel.velocity['vapor']['axial'][i-1], -10.0, 10.0)
            v_l_prev = np.clip(subchannel.velocity['liquid']['axial'][i-1], -10.0, 10.0)
            
            # 计算当前质量流量
            G_g = alpha * rho_g * v_g
            G_l = (1 - alpha) * rho_l * v_l
            
            # 计算上一个节点的质量流量
            G_g_prev = alpha_prev * rho_g * v_g_prev
            G_l_prev = (1 - alpha_prev) * rho_l * v_l_prev
            
            # 计算质量流量的变化率
            dG_g_dz = (G_g - G_g_prev) / self.grid.dz
            dG_l_dz = (G_l - G_l_prev) / self.grid.dz
            
            # 计算总的质量不平衡
            mass_imbalance = -(dG_g_dz + dG_l_dz)
            
            # 添加数值稳定项
            stability_factor = 1e-8
            mass_imbalance += stability_factor * (subchannel.pressure[i] - subchannel.previous_pressure[i])
            
            # 限制质量不平衡的大小
            max_imbalance = 1.0  # 最大允许的质量不平衡
            mass_imbalance = np.clip(mass_imbalance, -max_imbalance, max_imbalance)
            
            return mass_imbalance if np.isfinite(mass_imbalance) else 0.0
            
        except Exception as e:
            self.log_and_print(f"警告：节点 {i} 的质量不平衡计算失败: {str(e)}", level='warning')
            return 0.0

    def calculate_velocity_residual(self, subchannel):
        """计算速度残差
        
        Args:
            subchannel: 子通道对象
            
        Returns:
            float: 速度残差
        """
        n_axial = len(subchannel.alpha)
        velocity_residual = 0.0
        
        for i in range(1, n_axial-1):
            velocity_residual += np.abs(subchannel.velocity['vapor']['axial'][i] - 
                                      subchannel.previous_velocity['vapor']['axial'][i])
        
        return velocity_residual

    def calculate_pressure_residual(self, subchannel):
        """计算压力残差
        
        Args:
            subchannel: 子通道对象
            
        Returns:
            float: 压力残差
        """
        n_axial = len(subchannel.alpha)
        pressure_residual = 0.0
        
        for i in range(1, n_axial-1):
            pressure_residual += np.abs(subchannel.pressure[i] - 
                                      subchannel.previous_pressure[i])
        
        return pressure_residual

    def predict_velocity_piso(self, subchannel):
        """PISO算法的速度预测步骤"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算对流项（使用迎风格式）
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                v = subchannel.velocity[phase]['axial']
                rho = props['density']
                
                # 计算对流通量
                if v[i] > 0:
                    conv_flux = v[i] * (v[i] - v[i-1]) / self.grid.dz
                else:
                    conv_flux = v[i] * (v[i+1] - v[i]) / self.grid.dz
                
                # 计算压力梯度
                dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                
                # 预测速度（考虑对流项）
                subchannel.velocity[phase]['axial'][i] = (
                    subchannel.previous_velocity[phase]['axial'][i] -
                    self.dt * (conv_flux + dp_dz / rho)
                )

    def solve_pressure_correction_with_interfacial(self, subchannel):
        """考虑相间动量交换的压力修正方程求解"""
        n_axial = len(subchannel.alpha)
        
        try:
            # 构建压力修正方程系数矩阵
            A = np.zeros((n_axial, n_axial))
            b = np.zeros(n_axial)
            
            # 边界条件
            A[0, 0] = 1.0
            b[0] = 0.0  # 入口压力修正为0
            A[-1, -1] = 1.0
            b[-1] = 0.0  # 出口压力修正为0
            
            # 内部节点
            for i in range(1, n_axial-1):
                try:
                    # 获取物性
                    props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
                    props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
                    
                    # 确保密度为正值
                    rho_g = max(props_g['density'], 1e-6)
                    rho_l = max(props_l['density'], 1e-6)
                    
                    # 确保空泡份额在物理范围内
                    alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
                    
                    # 计算速度（添加限制）
                    v_g = np.clip(subchannel.velocity['vapor']['axial'][i], -10.0, 10.0)
                    v_l = np.clip(subchannel.velocity['liquid']['axial'][i], -10.0, 10.0)
                    v_r = abs(v_g - v_l)  # 相对速度
                    
                    # 计算相间动量交换系数（改进模型）
                    C_i = min(0.44 * alpha * (1 - alpha) * rho_l * v_r / self.grid.dz, 1e3)
                    
                    # 计算系数（考虑数值稳定性）
                    a_g = alpha / rho_g
                    a_l = (1 - alpha) / rho_l
                    
                    # 添加数值稳定项
                    stability_factor = 1e-6
                    
                    # 修改系数矩阵（考虑相间动量交换和数值稳定性）
                    A[i, i-1] = -(a_g + a_l) / (self.grid.dz ** 2)
                    A[i, i] = 2 * (a_g + a_l) / (self.grid.dz ** 2) + C_i + stability_factor
                    A[i, i+1] = -(a_g + a_l) / (self.grid.dz ** 2)
                    
                    # 计算源项（包含相间动量交换）
                    mass_imbalance = self.calculate_mass_imbalance(subchannel, i)
                    b[i] = mass_imbalance + C_i * (v_g - v_l)
                    
                except Exception as e:
                    self.log_and_print(f"警告：节点 {i} 的压力修正方程系数计算失败: {str(e)}", level='warning')
                    # 使用简单的拉普拉斯算子作为后备
                    A[i, i-1] = -1.0
                    A[i, i] = 2.0 + stability_factor
                    A[i, i+1] = -1.0
                    b[i] = 0.0
            
            # 增加对角优势以提高稳定性
            for i in range(n_axial):
                A[i,i] = max(A[i,i], 1.1 * np.sum(np.abs(A[i,:])))
            
            # 求解压力修正方程
            pressure_correction = np.linalg.solve(A, b)
            
            # 检查解的有效性
            if not np.all(np.isfinite(pressure_correction)):
                raise ValueError("压力修正解包含无效值")
            
            # 限制修正幅度
            max_correction = 0.1 * np.mean(np.abs(subchannel.pressure))
            pressure_correction = np.clip(pressure_correction, -max_correction, max_correction)
            
            # 应用压力修正（使用较小的松弛因子）
            omega = 0.1  # 更小的松弛因子
            subchannel.pressure += omega * pressure_correction
            
            # 确保压力保持在物理范围内
            min_pressure = 1e5  # 最小压力（1 bar）
            subchannel.pressure = np.maximum(subchannel.pressure, min_pressure)
            
        except Exception as e:
            self.log_and_print(f"警告：压力修正方程求解失败: {str(e)}", level='warning')
            # 使用简单的压力梯度估计
            dp = (subchannel.pressure[-1] - subchannel.pressure[0]) / (n_axial - 1)
            subchannel.pressure = np.linspace(subchannel.pressure[0], subchannel.pressure[-1], n_axial)

    def correct_velocity_with_interfacial(self, subchannel):
        """考虑相间动量交换的速度修正"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算压力梯度
            dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
            
            # 计算相间动量交换
            v_g = subchannel.velocity['vapor']['axial'][i]
            v_l = subchannel.velocity['liquid']['axial'][i]
            v_r = abs(v_g - v_l)
            
            alpha = subchannel.alpha[i]
            rho_l = props_l['density']
            
            # 相间动量交换系数
            C_i = 0.44 * alpha * (1 - alpha) * rho_l * v_r / self.grid.dz
            
            # 速度修正（考虑相间动量交换）
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                velocity_correction = (
                    -dp_dz / props['density'] +
                    C_i * (v_g - v_l) / props['density']
                )
                
                subchannel.velocity[phase]['axial'][i] += (
                    self.relaxation_factors['velocity'] * velocity_correction
                ) 

    def update_void_fraction_tvd(self, subchannel):
        """使用TVD格式更新空泡份额
        
        Args:
            subchannel: 子通道对象
        """
        n_axial = len(subchannel.alpha)
        alpha_new = np.zeros(n_axial)
        
        # 获取物性
        props_g = self.physics.calculate_properties(subchannel.pressure[0], PhaseType.VAPOR)
        props_l = self.physics.calculate_properties(subchannel.pressure[0], PhaseType.LIQUID)
        
        # 计算漂移速度
        v_gj, C0 = self.physics.calculate_drift_velocity(props_g['density'], 
                                                       props_l['density'])
        
        # 初始化漂移通量数组
        j_drift_array = np.zeros(n_axial)
        
        # 内部节点计算
        for i in range(1, n_axial-1):
            # 1. 计算TVD限制器
            r_plus = self.calculate_tvd_ratio(subchannel.alpha, i, 'forward')
            r_minus = self.calculate_tvd_ratio(subchannel.alpha, i, 'backward')
            
            # 使用van Leer限制器
            phi_plus = self.van_leer_limiter(r_plus)
            phi_minus = self.van_leer_limiter(r_minus)
            
            # 2. 计算通量
            # 2.1 计算混合速度
            j = (subchannel.alpha[i] * subchannel.velocity['vapor']['axial'][i] + 
                 (1-subchannel.alpha[i]) * subchannel.velocity['liquid']['axial'][i])
            
            # 2.2 计算实际气相速度（考虑漂移速度）
            v_g = C0 * j + v_gj
            
            # 2.3 计算漂移通量
            j_drift = subchannel.alpha[i] * (1-subchannel.alpha[i]) * v_gj
            j_drift_array[i] = j_drift
            
            # 3. 计算通量限制
            if v_g > 0:
                flux_plus = v_g * (subchannel.alpha[i] + 
                                  0.5 * phi_plus * (subchannel.alpha[i+1] - subchannel.alpha[i]))
                flux_minus = v_g * (subchannel.alpha[i-1] + 
                                   0.5 * phi_minus * (subchannel.alpha[i] - subchannel.alpha[i-1]))
            else:
                flux_plus = v_g * (subchannel.alpha[i+1] - 
                                  0.5 * phi_plus * (subchannel.alpha[i+1] - subchannel.alpha[i]))
                flux_minus = v_g * (subchannel.alpha[i] - 
                                   0.5 * phi_minus * (subchannel.alpha[i] - subchannel.alpha[i-1]))
            
            # 4. 计算漂移通量的导数
            if subchannel.previous_drift_flux is None:
                d_drift_dz = 0.0
            else:
                d_drift_dz = (j_drift - subchannel.previous_drift_flux[i]) / self.dt
            
            # 5. 更新空泡份额
            alpha_new[i] = subchannel.alpha[i] - self.dt/self.grid.dz * (flux_plus - flux_minus) + self.dt * d_drift_dz
        
        # 边界条件处理
        alpha_new[0] = subchannel.alpha[0]  # 入口边界条件
        alpha_new[-1] = alpha_new[-2]       # 出口外推
        
        # 物理限制
        alpha_new = np.clip(alpha_new, 0.0, 1.0)
            
        # 更新空泡份额
        subchannel.alpha = alpha_new
        
        # 保存当前漂移通量用于下一步计算
        subchannel.previous_drift_flux = j_drift_array

    def calculate_tvd_ratio(self, alpha: np.ndarray, i: int, direction: str) -> float:
        """计算TVD比率
        
        Args:
            alpha: 空泡份额数组
            i: 当前节点索引
            direction: 计算方向 ('forward' 或 'backward')
            
        Returns:
            float: TVD比率
        """
        eps = 1e-10  # 防止除零
        
        if direction == 'forward':
            if i < len(alpha)-2:
                numerator = alpha[i+2] - alpha[i+1]
                denominator = alpha[i+1] - alpha[i]
            else:
                numerator = 0.0
                denominator = 1.0
        else:  # backward
            if i > 1:
                numerator = alpha[i] - alpha[i-1]
                denominator = alpha[i-1] - alpha[i-2]
            else:
                numerator = 0.0
                denominator = 1.0
        
        return numerator / (denominator + eps)

    def van_leer_limiter(self, r: float) -> float:
        """van Leer限制器
        
        Args:
            r: TVD比率
            
        Returns:
            float: 限制器值
        """
        return (r + abs(r)) / (1 + abs(r))

    def minmod_limiter(self, r: float) -> float:
        """minmod限制器
        
        Args:
            r: TVD比率
            
        Returns:
            float: 限制器值
        """
        return max(0.0, min(1.0, r))

    def superbee_limiter(self, r: float) -> float:
        """superbee限制器
        
        Args:
            r: TVD比率
            
        Returns:
            float: 限制器值
        """
        return max(0.0, min(2.0*r, 1.0), min(r, 2.0)) 

    def solve_lateral_flow(self, subchannel, adjacent_channels, dt: float):
        """求解横向流动
        
        Args:
            subchannel: 当前子通道
            adjacent_channels: 相邻子通道列表
            dt: 时间步长
        """
        n_axial = len(subchannel.alpha)
        
        # 获取几何参数
        geom = subchannel.get_geometry_parameters()
        
        for i in range(1, n_axial-1):
            # 获取当前节点的物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 当前通道的参数
            alpha = subchannel.alpha[i]
            rho_g = props_g['density']
            rho_l = props_l['density']
            v_g = subchannel.velocity['vapor']['axial'][i]
            v_l = subchannel.velocity['liquid']['axial'][i]
            
            # 遍历相邻通道
            for adj_ch in adjacent_channels:
                # 获取连接参数
                conn_params = subchannel.get_connection_parameters(adj_ch)
                gap_width = conn_params['gap_width']
                gap_length = conn_params['gap_length']
                
                # 计算横向压力梯度
                dp_dy = (adj_ch.pressure[i] - subchannel.pressure[i]) / conn_params['centroid_distance']
                
                # 计算横向流动阻力系数（KIJ = 0.5）
                K_ij = 0.5
                
                # 计算横向流动速度（考虑FTM = 0.8）
                FTM = 0.8  # 横向动量因子
                v_g_lat = subchannel.velocity['vapor']['lateral'][i]
                v_l_lat = subchannel.velocity['liquid']['lateral'][i]
                
                # 计算横向流动阻力
                R_g = -K_ij * rho_g * v_g_lat * abs(v_g_lat) / gap_width
                R_l = -K_ij * rho_l * v_l_lat * abs(v_l_lat) / gap_width
                
                # 计算湍流混合项（β = 0.065）
                beta = 0.065  # 湍流交混系数
                w_t = self.physics.calculate_turbulent_mixing(
                    beta, gap_width, rho_g, rho_l, v_g, v_l
                )
                
                # 计算横向动量交换
                M_lat_g, M_lat_l = self.physics.calculate_lateral_momentum_exchange(
                    alpha, rho_g, rho_l, v_g, v_l,
                    adj_ch.velocity['vapor']['axial'][i],
                    adj_ch.velocity['liquid']['axial'][i],
                    gap_length, gap_width
                )
                
                # 计算空隙漂移效应
                v_gj, C0 = self.physics.calculate_drift_velocity(rho_g, rho_l)
                j_drift_lat = alpha * (1 - alpha) * v_gj * FTM
                
                # 更新横向速度
                # 气相
                dv_g_lat = dt * (
                    -dp_dy/rho_g +  # 压力梯度项
                    R_g/rho_g +     # 横向阻力项
                    M_lat_g/(rho_g * alpha * gap_width) +  # 横向动量交换项
                    w_t/(rho_g * alpha * gap_width)  # 湍流混合项
                )
                
                # 液相
                dv_l_lat = dt * (
                    -dp_dy/rho_l +  # 压力梯度项
                    R_l/rho_l +     # 横向阻力项
                    M_lat_l/(rho_l * (1-alpha) * gap_width) +  # 横向动量交换项
                    w_t/(rho_l * (1-alpha) * gap_width)  # 湍流混合项
                )
                
                # 应用松弛因子并更新横向速度
                subchannel.velocity['vapor']['lateral'][i] += (
                    self.relaxation_factors['velocity'] * dv_g_lat
                )
                subchannel.velocity['liquid']['lateral'][i] += (
                    self.relaxation_factors['velocity'] * dv_l_lat
                )
                
                # 更新质量通量（考虑横向流动的影响）
                G_g_lat = alpha * rho_g * subchannel.velocity['vapor']['lateral'][i]
                G_l_lat = (1-alpha) * rho_l * subchannel.velocity['liquid']['lateral'][i]
                
                # 保存当前时间步的横向质量通量
                subchannel.previous_mass_flux['vapor'][i] = G_g_lat
                subchannel.previous_mass_flux['liquid'][i] = G_l_lat

    def predict_velocity_piso(self, subchannel):
        """PISO算法的速度预测步骤"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算对流项（使用迎风格式）
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                v = subchannel.velocity[phase]['axial']
                rho = props['density']
                
                # 计算对流通量
                if v[i] > 0:
                    conv_flux = v[i] * (v[i] - v[i-1]) / self.grid.dz
                else:
                    conv_flux = v[i] * (v[i+1] - v[i]) / self.grid.dz
                
                # 计算压力梯度
                dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                
                # 预测速度（考虑对流项）
                subchannel.velocity[phase]['axial'][i] = (
                    subchannel.previous_velocity[phase]['axial'][i] -
                    self.dt * (conv_flux + dp_dz / rho)
                ) 